import sys
import time
import wave
import serial
import serial.tools.list_ports
from PySide6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
    QPushButton, QLabel, QComboBox, QLineEdit, QFileDialog, QMessageBox,
    QProgressBar
)
from PySide6.QtCore import Qt, QThread, Signal

# --- SerialWorker Class (remains the same as before) ---
class SerialWorker(QThread):
    """Handles serial communication in a separate thread."""
    status_update = Signal(str)
    progress_update = Signal(int) # Signal for progress percentage
    finished = Signal()
    error = Signal(str)

    def __init__(self, file_path, port, baud_rate, chunk_size=1024):
        super().__init__()
        self.file_path = file_path
        self.port = port
        self.baud_rate = baud_rate
        self.chunk_size = chunk_size
        self._is_running = True
        self.serial_connection = None
        self.wav_file = None

    def run(self):
        """Main worker task: open files/ports and stream data."""
        total_bytes_sent = 0
        total_file_size = 0

        try:
            # 1. Open WAV file
            self.status_update.emit(f"Opening WAV file: {self.file_path}")
            self.wav_file = wave.open(self.file_path, 'rb')

            # --- Get WAV file properties (optional but good practice) ---
            channels = self.wav_file.getnchannels()
            sample_width = self.wav_file.getsampwidth() # Bytes per sample
            frame_rate = self.wav_file.getframerate()
            num_frames = self.wav_file.getnframes()
            total_file_size = num_frames * channels * sample_width # Total PCM data size
            comp_type = self.wav_file.getcomptype()
            comp_name = self.wav_file.getcompname()

            if comp_type != 'NONE' or comp_name != 'not compressed':
                 raise ValueError("Only uncompressed PCM WAV files are supported.")

            self.status_update.emit(
                f"WAV Info: {channels}ch, {sample_width*8}-bit, {frame_rate}Hz, "
                f"~{total_file_size / 1024:.2f} KB PCM data"
            )
            # -------------------------------------------------------------

            # 2. Open Serial Port
            self.status_update.emit(f"Opening serial port {self.port} at {self.baud_rate} baud...")
            self.serial_connection = serial.Serial(
                port=self.port,
                baudrate=self.baud_rate,
                bytesize=serial.EIGHTBITS,
                parity=serial.PARITY_NONE,
                stopbits=serial.STOPBITS_ONE,
                timeout=1,  # Write timeout in seconds
                xonxoff=False, # No software flow control
                rtscts=False, # No hardware flow control (RTS/CTS)
                dsrdtr=False  # No hardware flow control (DSR/DTR)
            )
            self.status_update.emit(f"Serial port {self.port} opened.")

            # 3. Read and Send Data Chunks
            self.status_update.emit("Starting data transmission...")
            while self._is_running:
                data_chunk = self.wav_file.readframes(self.chunk_size)
                if not data_chunk:
                    break # End of file

                if not self._is_running: # Check again before writing
                    break

                try:
                    bytes_written = self.serial_connection.write(data_chunk)
                    self.serial_connection.flush() # Try to ensure data is sent physically
                    if bytes_written < len(data_chunk):
                        self.status_update.emit(f"Warning: Serial write timeout or incomplete write ({bytes_written}/{len(data_chunk)} bytes).")
                        # time.sleep(0.1) # Optional small delay on issues
                    total_bytes_sent += bytes_written

                    # Calculate and emit progress
                    if total_file_size > 0:
                        progress = int((total_bytes_sent / total_file_size) * 100)
                        self.progress_update.emit(progress)
                    else:
                         self.progress_update.emit(0)

                except serial.SerialTimeoutException:
                    self.status_update.emit("Error: Serial write timeout occurred. Stopping.")
                    self._is_running = False # Stop the loop
                    # self.error.emit("Serial write timeout.") # Uncomment if specific error needed
                    break
                except Exception as e:
                    raise RuntimeError(f"Error during serial write: {e}") # Propagate other write errors


            if self._is_running:
                self.status_update.emit("Transmission complete.")
                self.progress_update.emit(100) # Ensure progress reaches 100%
            else:
                self.status_update.emit("Transmission stopped by user.")

        except FileNotFoundError:
            self.error.emit(f"Error: File not found: {self.file_path}")
        except wave.Error as e:
            self.error.emit(f"Error reading WAV file: {e}. Is it a valid uncompressed PCM WAV?")
        except serial.SerialException as e:
            self.error.emit(f"Serial port error ({self.port}): {e}")
        except ValueError as e: # Catch our custom value error for compressed WAV
             self.error.emit(str(e))
        except Exception as e:
            self.error.emit(f"An unexpected error occurred: {e}")
        finally:
            # 4. Clean up
            if self.wav_file:
                self.wav_file.close()
                # self.status_update.emit("WAV file closed.") # Less verbose
            if self.serial_connection and self.serial_connection.is_open:
                try:
                    # time.sleep(0.1)
                    self.serial_connection.close()
                    # self.status_update.emit(f"Serial port {self.port} closed.") # Less verbose
                except serial.SerialException as e:
                    self.status_update.emit(f"Warning: Error closing serial port: {e}")

            self.finished.emit() # Signal completion regardless of success/failure

    def stop(self):
        """Signals the worker thread to stop."""
        self.status_update.emit("Stopping transmission...")
        self._is_running = False


# --- Main Application Window ---
class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Serial PCM Streamer (WAV)")
        self.setGeometry(100, 100, 550, 350)

        self.selected_file = None
        self.serial_worker = None
        self.is_streaming = False

        # --- Central Widget and Layout ---
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout(central_widget)

        # --- File Selection ---
        file_layout = QHBoxLayout()
        self.file_label = QLabel("Audio File:")
        self.file_path_edit = QLineEdit()
        self.file_path_edit.setReadOnly(True)
        self.browse_button = QPushButton("Browse...")
        self.browse_button.clicked.connect(self.browse_file)
        file_layout.addWidget(self.file_label)
        file_layout.addWidget(self.file_path_edit)
        file_layout.addWidget(self.browse_button)
        main_layout.addLayout(file_layout)

        # --- Serial Port Selection ---
        serial_layout = QHBoxLayout()
        self.port_label = QLabel("Serial Port:")
        self.port_combo = QComboBox()
        # Connect signal for when selection changes *after* initial population
        self.port_combo.currentIndexChanged.connect(self._update_start_button_state)
        self.refresh_ports_button = QPushButton("Refresh")
        self.refresh_ports_button.clicked.connect(self.refresh_ports)
        serial_layout.addWidget(self.port_label)
        serial_layout.addWidget(self.port_combo, 1)
        serial_layout.addWidget(self.refresh_ports_button)
        main_layout.addLayout(serial_layout)

        # --- Baud Rate Selection ---
        baud_layout = QHBoxLayout()
        self.baud_label = QLabel("Baud Rate:")
        self.baud_combo = QComboBox()
        common_baud_rates = [
            "9600", "19200", "38400", "57600", "115200",
            "230400", "460800", "921600", "1000000", "1500000", "2000000"
            ]
        self.baud_combo.addItems(common_baud_rates)
        self.baud_combo.setCurrentText("115200")
        baud_layout.addWidget(self.baud_label)
        baud_layout.addWidget(self.baud_combo)
        baud_layout.addStretch()
        main_layout.addLayout(baud_layout)

        # --- Control Buttons ---
        control_layout = QHBoxLayout()
        self.start_button = QPushButton("Start Streaming")
        self.start_button.clicked.connect(self.start_streaming)
        self.start_button.setEnabled(False) # Initially disabled
        self.stop_button = QPushButton("Stop Streaming")
        self.stop_button.clicked.connect(self.stop_streaming)
        self.stop_button.setEnabled(False)
        control_layout.addWidget(self.start_button)
        control_layout.addWidget(self.stop_button)
        main_layout.addLayout(control_layout)

        # --- Progress Bar ---
        self.progress_bar = QProgressBar()
        self.progress_bar.setRange(0, 100)
        self.progress_bar.setValue(0)
        self.progress_bar.setTextVisible(True)
        main_layout.addWidget(self.progress_bar)

        # --- Status Label ---
        self.status_label = QLabel("Status: Ready")
        self.status_label.setWordWrap(True)
        main_layout.addWidget(self.status_label)

        main_layout.addStretch()

        # --- Initial Population ---
        self.refresh_ports() # This will now also call _update_start_button_state

    def browse_file(self):
        """Opens a file dialog to select a WAV file."""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "Select WAV Audio File", "", "WAV Files (*.wav)"
        )
        if file_path:
            self.selected_file = file_path
            self.file_path_edit.setText(file_path)
            self.update_status(f"Selected file: {file_path}")
            self._update_start_button_state() # <--- Update button state here

    def refresh_ports(self):
        """Updates the serial port dropdown list."""
        # Temporarily block signals to prevent _update_start_button_state
        # firing multiple times during clear/add items
        self.port_combo.blockSignals(True)
        self.port_combo.clear()
        ports = serial.tools.list_ports.comports()
        if not ports:
            self.port_combo.addItem("No serial ports found")
            self.port_combo.setEnabled(False)
            # Keep start button disabled (will be handled by _update_start_button_state)
        else:
            for port in sorted(ports):
                self.port_combo.addItem(port.device)
            self.port_combo.setEnabled(True)
            # Enable start button conditionally (handled by _update_start_button_state)

        # Unblock signals *before* calling the update function
        self.port_combo.blockSignals(False)
        self._update_start_button_state() # <--- Update button state here
        self.update_status("Refreshed serial port list.")


    def _update_start_button_state(self):
        """Enables/disables the Start button based on current state."""
        #print(f"Debug: Updating start button state. Streaming: {self.is_streaming}, File: {self.selected_file}, Port: {self.port_combo.currentText()}")
        if self.is_streaming:
            self.start_button.setEnabled(False)
            return

        file_ok = self.selected_file is not None
        port_ok = False
        current_port_text = self.port_combo.currentText()
        if self.port_combo.isEnabled() and current_port_text and "No serial ports found" not in current_port_text:
            port_ok = True

        self.start_button.setEnabled(file_ok and port_ok)
        #print(f"Debug: Start button enabled: {file_ok and port_ok}")


    def start_streaming(self):
        """Starts the serial transmission process."""
        # Basic checks are still good practice here, though the button state
        # should prevent invalid calls if _update_start_button_state is correct.
        if self.is_streaming:
             # Should not happen if button logic is correct, but good failsafe
            QMessageBox.warning(self, "Warning", "Streaming is already in progress.")
            return

        # Redundant checks if button state is managed correctly, but safe:
        if not self.selected_file:
            QMessageBox.warning(self, "Warning", "Please select a WAV file first.")
            return
        selected_port = self.port_combo.currentText()
        if not selected_port or "No serial ports found" in selected_port:
            QMessageBox.warning(self, "Warning", "Please select a valid serial port.")
            return

        try:
            selected_baud = int(self.baud_combo.currentText())
        except ValueError:
             QMessageBox.critical(self, "Error", "Invalid baud rate selected.")
             return

        # --- Change GUI state to streaming ---
        self.set_gui_streaming_state(True) # This will also disable Start button via _update

        # --- Create and start worker ---
        self.serial_worker = SerialWorker(self.selected_file, selected_port, selected_baud)
        self.serial_worker.status_update.connect(self.update_status)
        self.serial_worker.progress_update.connect(self.update_progress)
        self.serial_worker.finished.connect(self.on_streaming_finished)
        self.serial_worker.error.connect(self.on_streaming_error)

        self.serial_worker.start()
        self.update_status(f"Attempting to start stream on {selected_port} at {selected_baud} baud...")
        self.progress_bar.setValue(0) # Reset progress bar


    def stop_streaming(self):
        """Stops the serial transmission process."""
        if self.serial_worker and self.is_streaming:
            self.serial_worker.stop()
            # GUI update will happen via the 'finished' or 'error' signal handlers
        else:
            self.update_status("Not currently streaming.")


    def update_status(self, message):
        """Updates the status label."""
        print(f"Status: {message}")
        self.status_label.setText(f"Status: {message}")

    def update_progress(self, value):
        """Updates the progress bar."""
        self.progress_bar.setValue(value)


    def set_gui_streaming_state(self, streaming):
        """Enables/disables GUI elements based on streaming state."""
        self.is_streaming = streaming
        # Controls directly affected by streaming state
        self.stop_button.setEnabled(streaming)
        self.browse_button.setEnabled(not streaming)
        self.port_combo.setEnabled(not streaming)
        self.refresh_ports_button.setEnabled(not streaming)
        self.baud_combo.setEnabled(not streaming)

        # Update start button based on *all* conditions, including the new streaming state
        self._update_start_button_state() # <--- Central place to update start button


    def on_streaming_finished(self):
        """Called when the worker thread finishes normally or is stopped."""
        self.update_status("Streaming finished or stopped.")
        self.set_gui_streaming_state(False) # Re-enable GUI, update start button
        self.serial_worker = None

    def on_streaming_error(self, error_message):
        """Called when the worker thread encounters an error."""
        self.update_status(f"Error: {error_message}")
        QMessageBox.critical(self, "Streaming Error", error_message)
        self.set_gui_streaming_state(False) # Re-enable GUI, update start button
        self.progress_bar.setValue(0)
        self.serial_worker = None


    def closeEvent(self, event):
        """Ensures the worker thread is stopped when closing the window."""
        if self.is_streaming and self.serial_worker:
            reply = QMessageBox.question(self, 'Confirm Exit',
                                         "Streaming is in progress. Stop streaming and exit?",
                                         QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No,
                                         QMessageBox.StandardButton.No)

            if reply == QMessageBox.StandardButton.Yes:
                self.update_status("Stopping stream due to application close...")
                self.serial_worker.stop()
                event.accept()
            else:
                event.ignore()
        else:
            event.accept()


# --- Application Entry Point ---
if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec())